from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import sklearn.decomposition as sk
from numpy.linalg import norm
from time import time
import matplotlib.pyplot as plt
from sys import stdout


def generate_dir_list(route="/Users/sinyeratlantis/Desktop/AR/"):
    dir_list = []
    for i in range(1, 121):
        for j in range(1, 15):
            temp = '00' + str(i) + '-0' + str(j)
            if i >= 10:
                temp = '0' + str(i) + '-0' + str(j)
            if i >= 100:
                temp = str(i) + '-0' + str(j)
            if j >= 10:
                temp = '00' + str(i) + '-' + str(j)
                if i >= 10:
                    temp = '0' + str(i) + '-' + str(j)
                if i >= 100:
                    temp = str(i) + '-' + str(j)
            img_dir = route + temp + '.bmp'
            dir_list.append(img_dir)
    return dir_list


def pca(flatten_imgs):  # 求主成分映射矩阵
    num_imgs, dim = flatten_imgs.shape  # 每行为一个样本，每列为一个维度
    mean_imgs = flatten_imgs.mean(axis=0)
    flatten_imgs = flatten_imgs - mean_imgs
    if dim <= num_imgs:  # 若单幅图片维度小于等于总图片数，使用奇异值分解SVD
        _, variance, projection = np.linalg.svd(flatten_imgs)
        projection = projection[:num_imgs]  # 取前nun_imgs维的数据
    else:  # 否则直接通过协方差矩阵求解
        covariance = np.dot(flatten_imgs, flatten_imgs.T)  # 计算协方差矩阵
        eigen, eigen_vector = np.linalg.eigh(covariance)  # 求解特征值和特征向量
        projection = np.dot(flatten_imgs.T, eigen_vector).T[::-1]  # 转换到特征向量空间得投影矩阵
    return projection, mean_imgs  # 返回投影矩阵、方差和均值


def plot(V, mean):  # 显示均值和前7个主成分图像
    img = np.array(Image.open(dir_list[0]))
    m, n = img.shape[0:2]
    plt.figure()
    plt.gray()
    plt.subplot(2, 4, 1)
    plt.imshow(mean.reshape(m, n))
    for i in range(7):
        plt.subplot(2, 4, i + 2)
        plt.imshow(V[i].reshape(m, n))
    plt.show()

dir_list = generate_dir_list()
imgs_matrix = np.array([np.array(Image.open(img)).flatten() for img in dir_list], 'f')
V, imgs_mean = pca(imgs_matrix)  # PCA处理

sk_pca = sk.PCA(7)
mean_imgs = imgs_matrix.mean(axis=0)
flatten_imgs = imgs_matrix - mean_imgs
sk_pca.fit(flatten_imgs.T)
V_pca = sk_pca.transform(imgs_matrix.T)

sk_nmf = sk.NMF(n_components=7, init='nndsvda', tol=5e-3)
sk_nmf.fit(imgs_matrix.T)
V_nmf = sk_nmf.transform(imgs_matrix.T)

plot(V, imgs_mean)
plot(V_pca.T, imgs_mean)
plot(V_nmf.T, imgs_mean)


def nmf(V, tole=0.001, t_limit=50, max_epoch=100):
    # 输入矩阵，近似量容忍度，时限，回合限
    W = np.random.randint(0, 9, size=(1680, 7))
    H = np.random.randint(0, 9, size=(7, 2000))
    t_init = time()

    w_grad = np.dot(W, np.dot(H, H.T)) - np.dot(V, H.T)
    h_grad = np.dot(np.dot(W.T, W), H) - np.dot(W.T, V)
    grad_init = norm(np.r_[w_grad, h_grad.T])  # np.concat ??

    w_tole = max(0.001, tole) * grad_init
    h_tole = w_tole

    for epoch in range(1, max_epoch):
        # 停止条件
        projnorm = norm(np.r_[w_grad[np.logical_or(w_grad < 0, W > 0)],
                              h_grad[np.logical_or(h_grad < 0, H > 0)]])
        if projnorm < tole * grad_init or time() - t_init > t_limit:
            break
        W, w_grad, iterW = nlssubprob(V.T, H.T, W.T, w_tole, 1000)
        W, w_grad = W.T, w_grad.T
        if iterW == 1:
            w_tole = 0.1 * w_tole
        (H, h_grad, iterH) = nlssubprob(V, W, H, h_tole, 1000)
        if iterH == 1:
            h_tole = 0.1 * h_tole

        if epoch % 10 == 0:
            stdout.write('.')

    print('epoch = %d Final proj-grad norm %f' % (epoch, projnorm))
    return W, H


def nlssubprob(V, W, h_init, tole, max_epoch):
    H = h_init
    WtV = np.dot(W.T, V)
    WtW = np.dot(W.T, W)
    alpha = 1
    beta = 0.1
    for epoch in range(1, max_epoch):
        grad = np.dot(WtW, H) - WtV
        projgrad = norm(grad[np.logical_or(grad < 0, H > 0)])
        if projgrad < tole:
            break
        for inner_iter in range(1, 20):
            Hn = H - alpha * grad
            Hn = np.where(Hn > 0, Hn, 0)
            d = Hn - H
            gradd = sum(grad * d)
            dQd = sum(np.dot(WtW, d) * d)
            suff_decr = 0.99 * gradd + 0.5 * dQd < 0
            if inner_iter == 1:
                print(suff_decr.shape)
                decr_alpha = not suff_decr.any()
                Hp = H
            if decr_alpha:
                if suff_decr.any():
                    H = Hn
                    break
                else:
                    alpha = alpha * beta
            else:
                if not suff_decr.any() or (Hp == Hn).all():
                    H = Hp
                    break
                else:
                    alpha = alpha / beta
                    Hp = Hn
    return H, grad, epoch

w = nmf(imgs_matrix)